import os
import sys
from typing import TYPE_CHECKING

from torch.utils.tensorboard import SummaryWriter
from zkl_metrics import MeanMetric, Metric
from zkl_training import ProcessingTaskPlugin

from .plugin_metrics import MetricsPlugin
from .utils import torch_distributed_get_info

if TYPE_CHECKING:
    from .training import GPTTraining


class SummaryPlugin(ProcessingTaskPlugin):
    def __init__(self):
        super().__init__()
        self.process_rank, _ = torch_distributed_get_info()
        self.last_summary_tokens_n: int | None = None
        self.summary_writer: SummaryWriter | None = None

    @property
    def task(self) -> 'GPTTraining':
        task = super().task
        from .training import GPTTraining
        assert isinstance(task, GPTTraining)
        return task

    @property
    def progress_tokens_n(self) -> int:
        return self.task.progress_tokens_n

    @property
    def metrics(self) -> dict[str, Metric]:
        return {
            attr_name[:-len('_metric')]: attr_value
            for attr_name, attr_value in vars(self).items()
            if attr_name.endswith('_metric') and isinstance(attr_value, Metric)}

    # task

    def on_before_run(self):
        self.last_summary_tokens_n = self.progress_tokens_n

        if self.process_rank == 0:
            self.summary_writer = SummaryWriter(log_dir=os.path.join(self.task.training_dir_path, "summaries"))

    def on_after_step(self):
        if self.check_need_summarize():
            self.summarize()
            self.last_summary_tokens_n = self.progress_tokens_n

    # summarizing

    def check_need_summarize(self):
        summary_tokens_n = self.task.training_hparams.summary_tokens_n
        if summary_tokens_n is None:
            return False

        if summary_tokens_n <= 0:
            return True

        warmup_tokens_n = self.task.training_hparams.warmup_tokens_n
        if warmup_tokens_n is not None:
            if self.last_summary_tokens_n < warmup_tokens_n:
                summary_tokens_n //= 10

        next_summary_tokens_n = ((self.last_summary_tokens_n // summary_tokens_n) + 1) * summary_tokens_n
        return self.progress_tokens_n >= next_summary_tokens_n

    def summarize(self):
        metrics = self.collect_metrics()
        if self.process_rank == 0:
            self.print_metrics(metrics)
            self.write_metrics(metrics)
        for metric in metrics.values():
            metric.reset()

    # collect

    def collect_metrics(self) -> dict[str, Metric]:
        metrics = {}

        for metric_name, metric_value in self.collect_metrics_from_training().items():
            metrics['train/' + metric_name] = metric_value

        for metric_name, metric_value in self.collect_metrics_from_validating().items():
            metrics['valid/' + metric_name] = metric_value

        return metrics

    def collect_metrics_from_training(self) -> dict[str, Metric]:
        metrics = {}
        for plugin in self.task.plugins:
            if isinstance(plugin, MetricsPlugin):
                metrics.update(plugin.metrics)
        return metrics

    def collect_metrics_from_validating(self) -> dict[str, Metric]:
        validating = self.task.validating_factory()
        return validating.run()

    # print & write

    def print_metrics(self, metrics: dict[str, Metric]):
        print(f"\nSummary (progress_tokens_n={self.progress_tokens_n})", file=sys.stderr)
        for name, metric in metrics.items():
            print(f"\t" + stringify_metric(name, metric), file=sys.stderr)

    def write_metrics(self, metrics: dict[str, Metric]):
        if self.summary_writer is not None:
            for name, metric in metrics.items():
                summarize_metric(self.summary_writer, name, metric, self.progress_tokens_n)


def stringify_metric(name: str, metric: Metric):
    if isinstance(metric, MeanMetric):
        value = metric.compute()
        if value is not None:
            if value.shape == ():
                value = value.detach().cpu().item()
                return f"{name}={value}"


def summarize_metric(writer: SummaryWriter, name: str, metric: Metric, step: int):
    if isinstance(metric, MeanMetric):
        value = metric.compute()
        if value is not None:
            if value.shape == ():
                value = value.detach().cpu().item()
                return writer.add_scalar(name, value, step)
